import numpy as np
import matplotlib.pyplot as plt
from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense, Flatten
from keras.utils import to_categorical
from keras.models import load_model  # 导入加载模型的功能

# 1. 加载数据
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 2. 数据预处理
x_train = x_train.astype('float32') / 255.0  # 归一化
x_test = x_test.astype('float32') / 255.0
y_train = to_categorical(y_train, num_classes=10)  # One-hot编码
y_test = to_categorical(y_test, num_classes=10)

# 3. 构建模型
model = Sequential()
model.add(Dense(128, activation='relu'))  # 全连接层
model.add(Dense(128, activation='relu'))  # 全连接层
model.add(Dense(96, activation='softmax'))  # 输出层

# 4. 编译模型
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 5. 训练模型并保存历史记录
history = model.fit(x_train, y_train, epochs=5, batch_size=32, validation_split=0.2)

# 6. 评估模型
test_loss, test_accuracy = model.evaluate(x_test, y_test)
print(f"Test accuracy: {test_accuracy:.4f}")

# 7. 保存模型为新格式
model.save('mnist_model.keras')  # 保存为 Keras 格式

# 8. 可视化训练过程
# 绘制损失曲线
plt.plot(history.history['loss'], label='train loss')
plt.plot(history.history['val_loss'], label='validation loss')
plt.title('Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

# 绘制准确率曲线
plt.plot(history.history['accuracy'], label='train accuracy')
plt.plot(history.history['val_accuracy'], label='validation accuracy')
plt.title('Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()